# -*- coding: utf-8 -*-

"""
@Author      : xxxxxxxxxxxxxxx
@Time        : 2024/6/12 14:18
"""
import os
import sys
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# 字体加载
from matplotlib import font_manager, rcParams

font_path = r"./font.ttf"
font_manager.fontManager.addfont(font_path)
prop = font_manager.FontProperties(fname=font_path)

# 字体设置
rcParams['font.family'] = 'sans-serif'  # 使用字体中的无衬线体
rcParams['font.sans-serif'] = prop.get_name()  # 根据名称设置字体
rcParams['font.size'] = 10  # 设置字体大小
rcParams['axes.unicode_minus'] = False  # 使坐标轴刻度标签正常显示正负号


def main_analysis(rir_file, obs_file, out_loc):
    df_obs = pd.read_csv(obs_file)
    df = pd.read_csv(rir_file)
    df = df[df['caseName'] == 'base']
    df['O3'] = df_obs['O₃(μg/m3)'].values[1:]
    df['O3小时监测浓度差值'] = df_obs['O₃(μg/m3)'].diff().values[1:]
    df['臭氧传输量T（O3）'] = df['O3小时监测浓度差值'] - df['NET(ppb/h)']
    df['时间'] = pd.to_datetime(df['BYMD(BJT)'], format="%Y-%m-%d %H:%M:%S")
    df.set_index('时间', inplace=True)
    df['hour'] = df.index.hour
    dft = df

    bottom = []
    for i in range(len(dft['臭氧传输量T（O3）'].values)):
        if (dft['臭氧传输量T（O3）'].values[i] >= 0) & (dft['NET(ppb/h)'].values[i] >= 0):
            bottom.append(dft['NET(ppb/h)'].values[i])
        elif (dft['臭氧传输量T（O3）'].values[i] < 0) & (dft['NET(ppb/h)'].values[i] < 0):
            bottom.append(dft['NET(ppb/h)'].values[i])
        else:
            bottom.append(0)

    width = 0.4
    fig = plt.figure()
    ax1 = fig.add_subplot(111)
    ax1.bar(range(len(dft['O3小时监测浓度差值'].values)), color='red', width=width, height=dft['NET(ppb/h)'].values, label='P(O$_{3}$)$_{net}$')
    ax1.plot(dft['O3小时监测浓度差值'].values, label='d(O$_{3}$)/dt', color='black', marker='o', markerfacecolor='red', markersize=4)
    ax1.bar(range(len(dft['O3小时监测浓度差值'].values)), color='green', width=width, bottom=bottom,
            height=dft['臭氧传输量T（O3）'].values,
            label='R(O$_{3}$)$_{trans}$')
    plt.legend(ncol=3, loc='upper left', fontsize=10)
    plt.margins()

    xticks = [f"{int(i)}时" for i in df['hour'].values[::4]]

    ybottom = -55
    ytop = 55
    ynum = 11
    plt.ylim(ybottom, ytop)
    plt.yticks(np.linspace(ybottom, ytop, ynum), ["{:.0f}".format(x) for x in np.linspace(ybottom, ytop, ynum)],
               fontsize=16, fontweight="bold")
    plt.hlines(y=0, ls='-', xmin=-100, xmax=len(dft['O3小时监测浓度差值']) * 2, colors='grey')
    plt.xlim(-0.5, len(dft['O3小时监测浓度差值']) - 0.5)
    plt.xticks(np.arange(0, len(dft['O3小时监测浓度差值']), 4), xticks, fontsize=16, fontweight="bold")
    plt.ylabel('浓度变化速率(ppb/h)', fontsize=16, fontweight="bold")

    ax2 = ax1.twinx()

    ax2.plot(dft['O3'].values, label='O$_{3}$', color='#029386', marker='o', markerfacecolor='y', markersize=4)
    plt.ylabel('浓度(μg/m${^3}$)', fontsize=16, fontweight="bold")
    plt.yticks(np.linspace(-200, 200, ynum), ["{:.0f}".format(x) for x in np.linspace(-200, 200, ynum)], fontsize=16, fontweight="bold")
    plt.legend(fontsize=10)

    os.makedirs(out_loc, exist_ok=True)
    plt.savefig(os.path.join(out_loc, '收支分析.png'), bbox_inches='tight', dpi=600)
    # plt.show()


if __name__ == '__main__':
    if len(sys.argv) < 4:
        print(f"[Error]: Run with python {sys.argv[0]} './5_rir.csv' './5.csv' './picture'")
    else:
        main_analysis(sys.argv[1], sys.argv[2], sys.argv[3])
